{
  "nbformat": 4,
  "nbformat_minor": 0,
  "metadata": {
    "colab": {
      "name": "function.ipynb",
      "version": "0.3.2",
      "provenance": [],
      "private_outputs": true,
      "collapsed_sections": [
        "Jxv6goXm7oGF"
      ],
      "toc_visible": true
    },
    "kernelspec": {
      "display_name": "Python 3",
      "name": "python3"
    }
  },
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {
        "colab_type": "text",
        "id": "Jxv6goXm7oGF"
      },
      "source": [
        "##### Copyright 2018 The TensorFlow Authors.\n",
        "\n",
        "Licensed under the Apache License, Version 2.0 (the \"License\");"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "colab_type": "code",
        "id": "llMNufAK7nfK",
        "colab": {}
      },
      "source": [
        "#@title Licensed under the Apache License, Version 2.0 (the \"License\"); { display-mode: \"form\" }\n",
        "# you may not use this file except in compliance with the License.\n",
        "# You may obtain a copy of the License at\n",
        "#\n",
        "# https://www.apache.org/licenses/LICENSE-2.0\n",
        "#\n",
        "# Unless required by applicable law or agreed to in writing, software\n",
        "# distributed under the License is distributed on an \"AS IS\" BASIS,\n",
        "# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n",
        "# See the License for the specific language governing permissions and\n",
        "# limitations under the License."
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "colab_type": "text",
        "id": "8Byow2J6LaPl"
      },
      "source": [
        "# TensorFlow 2.0 での tf.function と AutoGraph"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "colab_type": "text",
        "id": "kGXS3UWBBNoc"
      },
      "source": [
        "<table class=\"tfo-notebook-buttons\" align=\"left\">\n",
        "  <td>\n",
        "    <a target=\"_blank\" href=\"https://www.tensorflow.org/guide/function\"><img src=\"https://www.tensorflow.org/images/tf_logo_32px.png\" />View on TensorFlow.org</a>\n",
        "  </td>\n",
        "  <td>\n",
        "    <a target=\"_blank\" href=\"https://colab.research.google.com/github/tensorflow/docs/blob/master/site/ja/guide/function.ipynb\"><img src=\"https://www.tensorflow.org/images/colab_logo_32px.png\" />Run in Google Colab</a>\n",
        "  </td>\n",
        "  <td>\n",
        "    <a target=\"_blank\" href=\"https://github.com/tensorflow/docs/blob/master/site/ja/guide/function.ipynb\"><img src=\"https://www.tensorflow.org/images/GitHub-Mark-32px.png\" />View source on GitHub</a>\n",
        "  </td>\n",
        "</table>"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "TPrI2hpDnNgG",
        "colab_type": "text"
      },
      "source": [
        "Note: これらのドキュメントは私たちTensorFlowコミュニティが翻訳したものです。コミュニティによる 翻訳は**ベストエフォート**であるため、この翻訳が正確であることや[英語の公式ドキュメント](https://www.tensorflow.org/?hl=en)の 最新の状態を反映したものであることを保証することはできません。 この翻訳の品質を向上させるためのご意見をお持ちの方は、GitHubリポジトリ[tensorflow/docs](https://github.com/tensorflow/docs)にプルリクエストをお送りください。 コミュニティによる翻訳やレビューに参加していただける方は、 [docs-ja@tensorflow.org メーリングリスト](https://groups.google.com/a/tensorflow.org/forum/#!forum/docs-ja)にご連絡ください。"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "colab_type": "text",
        "id": "CydFK2CL7ZHA"
      },
      "source": [
        "TensorFlow 2.0 では Eager Execution の使いやすさとTensorFlow 1.0 のパワーとを同時に提供します。この統合の中核となるのは `tf.function` です。これは Python の構文のサブセットを移植可能でハイパフォーマンスな TensorFlow のグラフに変換します。\n",
        "\n",
        "`tf.function`の魅力的な特徴に AutoGraph があります。これはグラフを Python の構文そのものを用いて記述できるようにします。\n",
        "AutoGraph で利用可能な Python の機能の一覧は、[AutoGraph Capabilities and Limitations (Autograph の性能と制限事項)](https://github.com/tensorflow/tensorflow/blob/master/tensorflow/python/autograph/g3doc/reference/limitations.md) で確認できます。また、`tf.function`の詳細については RFC [TF 2.0: Functions, not Sessions](https://github.com/tensorflow/community/blob/master/rfcs/20180918-functions-not-sessions-20.md) を参照してください。AutoGraph の詳細については `tf.autograph` を参照してください。\n",
        "\n",
        "このチュートリアルでは `tf.function` と AutoGraph の基本的な特徴についてひととおり確認します。"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "colab_type": "text",
        "id": "n4EKOpw9mObL"
      },
      "source": [
        "## セットアップ\n",
        "\n",
        "TensorFlow 2.0 をインポートして、TF 2.0 モードを有効にしてください。"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "colab_type": "code",
        "id": "V9oECvVSI1Kj",
        "colab": {}
      },
      "source": [
        "from __future__ import absolute_import, division, print_function, unicode_literals\n",
        "import numpy as np"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "colab_type": "code",
        "id": "mT7meGqrZTz9",
        "colab": {}
      },
      "source": [
        "try:\n",
        "  !pip install tf-nightly-2.0-preview\n",
        "except Exception:\n",
        "  pass\n",
        "import tensorflow as tf"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "colab_type": "text",
        "id": "77AsVr1GGtBP"
      },
      "source": [
        "## `tf.function` デコレータ\n",
        "\n",
        "`tf.function`を用いてある関数にアノテーションを付けたとしても、一般の関数と変わらずに呼び出せます。一方、実行時にはその関数はグラフへとコンパイルされます。これにより、より高速な実行や、 GPU や TPU での実行、SavedModel へのエクスポートといった利点が得られます。"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "colab_type": "code",
        "id": "FhIg7-z6HNWj",
        "colab": {}
      },
      "source": [
        "@tf.function\n",
        "def simple_nn_layer(x, y):\n",
        "  return tf.nn.relu(tf.matmul(x, y))\n",
        "\n",
        "\n",
        "x = tf.random.uniform((3, 3))\n",
        "y = tf.random.uniform((3, 3))\n",
        "\n",
        "simple_nn_layer(x, y)"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "colab_type": "text",
        "id": "U-LAE4pMNR9g"
      },
      "source": [
        "アノテーションの結果を調べてみると、 TensorFlow ランタイムとのやり取りのすべてを処理する特別な呼び出し可能オブジェクトを確認できます。 "
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "colab_type": "code",
        "id": "q4t2iuS7Nqc0",
        "colab": {}
      },
      "source": [
        "simple_nn_layer"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "colab_type": "text",
        "id": "DqeefLGNXjZQ"
      },
      "source": [
        "記述したコードで複数の関数を利用していたとしても、すべての関数にアノテーションを付ける必要はありません。アノテーションをつけた関数から呼び出されるすべての関数は、グラフモードで実行されます。"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "colab_type": "code",
        "id": "3VGF7tlVXiZY",
        "colab": {}
      },
      "source": [
        "def linear_layer(x):\n",
        "  return 2 * x + 1\n",
        "\n",
        "\n",
        "@tf.function\n",
        "def deep_net(x):\n",
        "  return tf.nn.relu(linear_layer(x))\n",
        "\n",
        "\n",
        "deep_net(tf.constant((1, 2, 3)))"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "colab_type": "text",
        "id": "yQvg6ZSKWyqE"
      },
      "source": [
        "グラフが大量の軽量な演算から構成される場合、関数は Eager Execution で実行するコードよりも高速になる場合があります。しかし、 graph が少量の (畳み込み演算のような) 計算に時間のかかる演算からなる場合、高速化はそれほど見込めないでしょう。"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "colab_type": "code",
        "id": "0EL6lVwEWuFo",
        "colab": {}
      },
      "source": [
        "import timeit\n",
        "conv_layer = tf.keras.layers.Conv2D(100, 3)\n",
        "\n",
        "@tf.function\n",
        "def conv_fn(image):\n",
        "  return conv_layer(image)\n",
        "\n",
        "image = tf.zeros([1, 200, 200, 100])\n",
        "# warm up\n",
        "conv_layer(image); conv_fn(image)\n",
        "print(\"Eager conv:\", timeit.timeit(lambda: conv_layer(image), number=10))\n",
        "print(\"Function conv:\", timeit.timeit(lambda: conv_fn(image), number=10))\n",
        "print(\"Note how there's not much difference in performance for convolutions\")\n"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "colab_type": "code",
        "id": "L4zj-jpH0jKH",
        "colab": {}
      },
      "source": [
        "lstm_cell = tf.keras.layers.LSTMCell(10)\n",
        "\n",
        "@tf.function\n",
        "def lstm_fn(input, state):\n",
        "  return lstm_cell(input, state)\n",
        "\n",
        "input = tf.zeros([10, 10])\n",
        "state = [tf.zeros([10, 10])] * 2\n",
        "# warm up\n",
        "lstm_cell(input, state); lstm_fn(input, state)\n",
        "print(\"eager lstm:\", timeit.timeit(lambda: lstm_cell(input, state), number=10))\n",
        "print(\"function lstm:\", timeit.timeit(lambda: lstm_fn(input, state), number=10))\n"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "colab_type": "text",
        "id": "ohbSnA79mcJV"
      },
      "source": [
        "## Python の制御フローの利用\n",
        "\n",
        "`tf.function`の内部でデータに依存した制御フローを用いる場合、Pythonの制御フロー構文を用いることができます。AutoGraph はそれらの構文を TensorFlow の Ops に書き換えます。たとえば、 `Tensor` に依存する `if` 文は、`tf.cond()` に変換されます。\n",
        "\n",
        "次の例では `x` は `Tensor` です。ですが、 `if` 文は期待するどおりに動作しています。"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "colab_type": "code",
        "id": "aA3gOodCBkOw",
        "colab": {}
      },
      "source": [
        "@tf.function\n",
        "def square_if_positive(x):\n",
        "  if x > 0:\n",
        "    x = x * x\n",
        "  else:\n",
        "    x = 0\n",
        "  return x\n",
        "\n",
        "\n",
        "print('square_if_positive(2) = {}'.format(square_if_positive(tf.constant(2))))\n",
        "print('square_if_positive(-2) = {}'.format(square_if_positive(tf.constant(-2))))"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "colab_type": "text",
        "id": "GMiCUkdyoq98"
      },
      "source": [
        "Note: この例ではスカラー値を用いた単純な条件を用いています。実利用する場合、典型的には<a href=\"#batching\">バッチ処理</a>が用いられます。"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "colab_type": "text",
        "id": "m-jWmsCmByyw"
      },
      "source": [
        "AutoGraph は `while`, `for`, `if`, `break`, `continue`, `return` といった典型的なPythonの構文をサポートしています。また、これらを入れ子にして利用する場合もサポートしています。つまり、`Tensor`を返す式を`while` 文や `if` 文の条件式として用いることが可能です。また、`for` 文で `Tensor` の要素に渡って反復することも可能です。"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "colab_type": "code",
        "id": "toxKBOXbB1ro",
        "colab": {}
      },
      "source": [
        "@tf.function\n",
        "def sum_even(items):\n",
        "  s = 0\n",
        "  for c in items:\n",
        "    if c % 2 > 0:\n",
        "      print(c)\n",
        "      continue\n",
        "    s += c\n",
        "  return s\n",
        "\n",
        "\n",
        "sum_even(tf.constant([10, 12, 15, 20]))"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "colab_type": "text",
        "id": "AtDaLrbySw4j"
      },
      "source": [
        "より高度な使い方をするユーザーのために、AutoGraph は低レベルAPIも提供しています。次の例では AutoGraph が生成したコードを確認できます。"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "colab_type": "code",
        "id": "aRsde3x_SjTQ",
        "colab": {}
      },
      "source": [
        "print(tf.autograph.to_code(sum_even.python_function))"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "colab_type": "text",
        "id": "rvJXCfk8VkLf"
      },
      "source": [
        "次はより複雑な制御フローの例です。"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "colab_type": "code",
        "id": "h-Z87IJqVlKl",
        "colab": {}
      },
      "source": [
        "@tf.function\n",
        "def fizzbuzz(n):\n",
        "  msg = tf.constant('')\n",
        "  for i in tf.range(n):\n",
        "    if i % 3 == 0:\n",
        "      tf.print('Fizz')\n",
        "    elif i % 5 == 0:\n",
        "      tf.print('Buzz')\n",
        "    else:\n",
        "      tf.print(i)\n",
        "\n",
        "fizzbuzz(tf.constant(15))"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "colab_type": "text",
        "id": "h_Y4uC1R1B55"
      },
      "source": [
        "## Keras での AutoGraph の利用\n",
        "\n",
        "`tf.function` はオブジェクトのメソッドに対しても利用できます。たとえば、カスタムしたKeras モデルにデコレーターを適用できます、典型的には `call` 関数にアノテーションを付けることで実現できるでしょう。より詳細が必要な場合、`tf.keras` を確認してください。"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "colab_type": "code",
        "id": "cR6mpLKP1HLe",
        "colab": {}
      },
      "source": [
        "class CustomModel(tf.keras.models.Model):\n",
        "\n",
        "  @tf.function\n",
        "  def call(self, input_data):\n",
        "    if tf.reduce_mean(input_data) > 0:\n",
        "      return input_data\n",
        "    else:\n",
        "      return input_data // 2\n",
        "\n",
        "\n",
        "model = CustomModel()\n",
        "\n",
        "model(tf.constant([-2, -4]))"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "colab_type": "text",
        "id": "NTEvpBK9f8kj"
      },
      "source": [
        "## 副作用\n",
        "\n",
        "Eager モードのように、通常の場合 `tf.function` の中で、`tf.assign` や `tf.print` といった副作用のある命令を実行できます。また、実行時の順序を保つために、処理順について必要な依存関係を書き加えます。"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "colab_type": "code",
        "id": "-Wd6i8S9gcuC",
        "colab": {}
      },
      "source": [
        "v = tf.Variable(5)\n",
        "\n",
        "@tf.function\n",
        "def find_next_odd():\n",
        "  v.assign(v + 1)\n",
        "  if v % 2 == 0:\n",
        "    v.assign(v + 1)\n",
        "\n",
        "\n",
        "find_next_odd()\n",
        "v"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "colab_type": "text",
        "id": "4LfnJjm0Bm0B"
      },
      "source": [
        "## 例: シンプルなモデルの学習\n",
        "\n",
        "AutoGraph はこれまで見てきたよりもずっと多くの演算を TensorFlow の内部で実行できます。たとえば、学習のためのループ処理は単に制御フローなので、実際にそれを TensorFlow に持ち込んで処理できます。"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "colab_type": "text",
        "id": "Em5dzSUOtLRP"
      },
      "source": [
        "### データのダウンロード"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "colab_type": "code",
        "id": "xqoxumv0ssQW",
        "colab": {}
      },
      "source": [
        "def prepare_mnist_features_and_labels(x, y):\n",
        "  x = tf.cast(x, tf.float32) / 255.0\n",
        "  y = tf.cast(y, tf.int64)\n",
        "  return x, y\n",
        "\n",
        "def mnist_dataset():\n",
        "  (x, y), _ = tf.keras.datasets.mnist.load_data()\n",
        "  ds = tf.data.Dataset.from_tensor_slices((x, y))\n",
        "  ds = ds.map(prepare_mnist_features_and_labels)\n",
        "  ds = ds.take(20000).shuffle(20000).batch(100)\n",
        "  return ds\n",
        "\n",
        "train_dataset = mnist_dataset()"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "colab_type": "text",
        "id": "znmy4l8ntMvW"
      },
      "source": [
        "### モデルの定義"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "colab_type": "code",
        "id": "ltxyJVWTqNAO",
        "colab": {}
      },
      "source": [
        "model = tf.keras.Sequential((\n",
        "    tf.keras.layers.Reshape(target_shape=(28 * 28,), input_shape=(28, 28)),\n",
        "    tf.keras.layers.Dense(100, activation='relu'),\n",
        "    tf.keras.layers.Dense(100, activation='relu'),\n",
        "    tf.keras.layers.Dense(10)))\n",
        "model.build()\n",
        "optimizer = tf.keras.optimizers.Adam()"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "colab_type": "text",
        "id": "oeYV6mKnJGMr"
      },
      "source": [
        "### 学習のためのループ処理の定義"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "colab_type": "code",
        "id": "3xtg_MMhJETd",
        "colab": {}
      },
      "source": [
        "compute_loss = tf.keras.losses.SparseCategoricalCrossentropy(from_logits=True)\n",
        "\n",
        "compute_accuracy = tf.keras.metrics.SparseCategoricalAccuracy()\n",
        "\n",
        "\n",
        "def train_one_step(model, optimizer, x, y):\n",
        "  with tf.GradientTape() as tape:\n",
        "    logits = model(x)\n",
        "    loss = compute_loss(y, logits)\n",
        "\n",
        "  grads = tape.gradient(loss, model.trainable_variables)\n",
        "  optimizer.apply_gradients(zip(grads, model.trainable_variables))\n",
        "\n",
        "  compute_accuracy(y, logits)\n",
        "  return loss\n",
        "\n",
        "\n",
        "@tf.function\n",
        "def train(model, optimizer):\n",
        "  train_ds = mnist_dataset()\n",
        "  step = 0\n",
        "  loss = 0.0\n",
        "  accuracy = 0.0\n",
        "  for x, y in train_ds:\n",
        "    step += 1\n",
        "    loss = train_one_step(model, optimizer, x, y)\n",
        "    if step % 10 == 0:\n",
        "      tf.print('Step', step, ': loss', loss, '; accuracy', compute_accuracy.result())\n",
        "  return step, loss, accuracy\n",
        "\n",
        "step, loss, accuracy = train(model, optimizer)\n",
        "print('Final step', step, ': loss', loss, '; accuracy', compute_accuracy.result())"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "colab_type": "text",
        "id": "SnsumiP6eRYL"
      },
      "source": [
        "## バッチ処理\n",
        "\n",
        "実際のアプリケーションにおいて、処理をバッチにまとめることはパフォーマンスの観点から重要です。AutoGraphを用いるのにもっとも適しているコードは、制御フローを _バッチ_ の単位で決定するようなコードです。もし、個々の _要素_ の単位で制御を決定する場合、パフォーマンスを保つために batch API を試してみてください。\n",
        "\n",
        "一例として、次の Python コードががあったとします。"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "colab_type": "code",
        "id": "t31QoERiNccJ",
        "colab": {}
      },
      "source": [
        "def square_if_positive(x):\n",
        "  return [i ** 2 if i > 0 else i for i in x]\n",
        "\n",
        "\n",
        "square_if_positive(range(-5, 5))"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "colab_type": "text",
        "id": "kSeEJ76uNgwD"
      },
      "source": [
        "TensorFlowに同等の処理を行わせる場合、次のように記述したくなるかもしれません。 (これは実際には動作します!)"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "colab_type": "code",
        "id": "RqR8WzSzNf87",
        "colab": {}
      },
      "source": [
        "@tf.function\n",
        "def square_if_positive_naive(x):\n",
        "  result = tf.TensorArray(tf.int32, size=x.shape[0])\n",
        "  for i in tf.range(x.shape[0]):\n",
        "    if x[i] > 0:\n",
        "      result = result.write(i, x[i] ** 2)\n",
        "    else:\n",
        "      result = result.write(i, x[i])\n",
        "  return result.stack()\n",
        "\n",
        "\n",
        "square_if_positive_naive(tf.range(-5, 5))"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "colab_type": "text",
        "id": "gTcyWXVGN3gS"
      },
      "source": [
        "しかし、この場合、次のように書くこともできます。"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "colab_type": "code",
        "id": "VO2f6x-lNfVj",
        "colab": {}
      },
      "source": [
        "def square_if_positive_vectorized(x):\n",
        "  return tf.where(x > 0, x ** 2, x)\n",
        "\n",
        "\n",
        "square_if_positive_vectorized(tf.range(-5, 5))"
      ],
      "execution_count": 0,
      "outputs": []
    }
  ]
}
